import logging
import os
import time
from collections import defaultdict
from pathlib import Path

import cv2
import imageio
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from plyfile import PlyData, PlyElement
from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import make_grid, save_image
from tqdm import tqdm

from dataset.rgbd_dataset import ShapeViewDataset, Voxelizer
from network.loss import SRNLossCriterion
from network.radiance_field import RaidanceField
from network.renderer import VolumeRenderer
from options import config_parser, config_parser_extract_mesh
from utils import data_utils, exp_util, motion_util
from rgbd_nerf_train import RGBD_NeRF

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    parser = config_parser_extract_mesh()
    args = parser.parse_args()
    logging.info(args)
    device = torch.device("cuda:0")
    # writer = SummaryWriter(os.path.join(args.log_path, args.exp_name,"tensorboard"))

    ### init data & model
    views = 1 ## current we only consider views eqs 1
    view_ids = [0, 30, 60, 90] ## specify views for recon
    assert len(view_ids) == args.num_view
    rgbd_dataset = ShapeViewDataset(args.root_path, views, args.num_view, load_depth=True, resolution=args.reso)
    vox = Voxelizer(args, device)

    samples = rgbd_dataset.get_base_sequence(index = 0, view_ids=view_ids)
    data_seq, encoder_states = vox.preprocess(samples)

    trainer = RGBD_NeRF(args)
    if args.resume:
        if args.load_path is not None:
            trainer.load_state_dict(torch.load(args.load_path), strict=True)
        else:
            
            assert os.path.isdir(os.path.join(args.checkpoint_path, args.exp_name))
            load_path = os.path.join(args.checkpoint_path, args.exp_name, "latest.pth.tar")
            logging.info(f"ckpt is not specified, use the latest file {load_path} instead.")
            trainer.load_state_dict(torch.load(load_path), strict=True)
    trainer.to(device)

    criterion = SRNLossCriterion(args)

    ## test marching cubes lewiner
    with torch.no_grad():
        trainer.eval()
        trainer.extract_mesh(encoder_states, args.mc_threshold, args.mc_num_samples_per_halfvoxel)


    ## test given sequence
    with torch.no_grad():
        trainer.eval()
        for i in range(90, 120):

            sample = rgbd_dataset[i]
            array_name = ["colors", "uv", "alpha", 'extrinsics','intrinsics', 'depths']
            if isinstance(sample['colors'], np.ndarray):
                transfer = lambda x: torch.from_numpy(x).to(device)
                for j in array_name:
                    sample[j]=transfer(sample[j])

            trainer.valid_step(sample, encoder_states, criterion, step=i)

